package com.hob.openai.utils;

import com.theokanning.openai.OpenAiService;
import com.theokanning.openai.completion.CompletionChoice;
import com.theokanning.openai.completion.CompletionRequest;

import java.util.Arrays;
import java.util.List;

public class SqlAiUtils {

    /**
     * 生成需求
     * @return
     */
    public static List<CompletionChoice> generatorSqlQuery(String tables, String demand){
        String keyword = "";
        if(demand.contains("查询")||demand.contains("select")){
            keyword = "SELECT";
        }else if(demand.contains("删除数据")||demand.contains("delete")){
            keyword = "DELETE";
        }else if(demand.contains("删除")||demand.contains("drop")){
            keyword = "DROP";
        }else if(demand.contains("修改数据")||demand.contains("update")){
            keyword = "UPDATE";
        }else if(demand.contains("修改")||demand.contains("alter")){
            keyword = "ALTER";
        }else if(demand.contains("增加")||demand.contains("insert")){
            keyword = "INSERT";
        }else if(demand.contains("创建")||demand.contains("create")){
            keyword = "CREATE";
        }else if(demand.contains("显示")||demand.contains("show")){
            keyword = "SHOW";
        }else if(demand.contains("清空")||demand.contains("truncate")){
            keyword = "TRUNCATE";
        }
        OpenAiService service = new OpenAiService(TOKEN);
        CompletionRequest completionRequest = CompletionRequest.builder()
                .model("code-davinci-002")
                .prompt(String.format("%s %s\\n%s",tables,demand,keyword))
                .temperature(0.0D)
                .maxTokens(1000)
                .topP(1D)
                .frequencyPenalty(0D)
                .presencePenalty(0D)
                .stop(Arrays.asList("#",";"))
                .build();
        List<CompletionChoice> choices = service.createCompletion(completionRequest).getChoices();
        for (CompletionChoice choice : choices) {
            choice.setText(keyword+choice.getText());
        }
        return choices;
    }

}
